import numpy as np
import open3d as o3d
from sklearn.cluster import DBSCAN
import matplotlib.pyplot as plt

class SegmentationMethod:
    """分割方法枚举（更新方法名称）"""
    EUCLIDEAN = "欧几里得聚类"  # 替换原 DBSCAN
    CURVATURE = "法线曲率分割"
    @staticmethod
    def get_all_methods():
        """获取所有支持的分割方法"""
        return [
            SegmentationMethod.EUCLIDEAN,
            SegmentationMethod.CURVATURE
        ]

class PointCloudSegmentor:
    """点云分割工具类（集中所有分割方法）"""
    @staticmethod
    def segment(point_cloud_data, method: str):
        """
        统一分割接口
        :param point_cloud_data: (M,3) 点云数据
        :param method: 分割方法（来自SegmentationMethod）
        :return: 带颜色的Open3D点云对象、分割标签
        """
        if method == SegmentationMethod.EUCLIDEAN:
            return PointCloudSegmentor.euclidean_segment(point_cloud_data)
        elif method == SegmentationMethod.CURVATURE:
            return PointCloudSegmentor.curvature_segment(point_cloud_data)
        else:
            raise ValueError(f"不支持的分割方法：{method}")

    @staticmethod
    def euclidean_segment(point_cloud_data):
        """欧几里得聚类分割（替换原有DBSCAN，新增去噪和鲜艳配色）"""
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(point_cloud_data)
        
        # 预处理：去噪（统计离群点去除）
        pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
        # 注意：去噪后点云数量会减少，更新点云数据（避免标签与点云不匹配）
        denoised_points = np.asarray(pcd.points)
        print(f"去噪前点数：{len(point_cloud_data)}，去噪后点数：{len(denoised_points)}")
        
        # 欧几里得聚类（基于DBSCAN）
        labels = np.array(pcd.cluster_dbscan(eps=0.02, min_points=50, print_progress=True))
        max_label = labels.max()
        
        # 为不同部件分配鲜艳颜色（使用Set3配色，支持更多部件）
        colors = plt.cm.Set3(np.linspace(0, 1, max_label + 1))[:, :3]
        # 噪声点（标签-1）设为灰色（保持与曲率分割一致）
        colored_labels = labels.copy()
        noise_mask = colored_labels == -1
        if np.any(noise_mask):
            colored_labels[noise_mask] = max_label + 1
            colors = np.vstack([colors, [0.5, 0.5, 0.5]])  # 灰色
        
        # 设置点云颜色
        pcd.colors = o3d.utility.Vector3dVector(colors[colored_labels])
        
        # 统计信息
        part_count = max_label + 1 if max_label >= 0 else 0
        noise_count = np.sum(noise_mask)
        print(f"欧几里得聚类分割出 {part_count} 个部件，噪声点 {noise_count} 个")
        
        # 仅返回系统需要的 pcd 和 labels（clusters 如需使用可后续扩展）
        return pcd, labels

    @staticmethod
    def curvature_segment(point_cloud_data):
        """法线曲率分割方法（保持不变）"""
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(point_cloud_data)
        
        # 估计法线
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
        
        # 计算曲率（通过法线变化）
        points = np.asarray(pcd.points)
        normals = np.asarray(pcd.normals)
        
        # 使用KNN查找邻近点计算局部曲率
        pcd_tree = o3d.geometry.KDTreeFlann(pcd)
        curvatures = []
        
        for i in range(len(points)):
            [k, idx, _] = pcd_tree.search_knn_vector_3d(pcd.points[i], 30)
            # 计算法线变化作为曲率估计
            neighbor_normals = normals[idx]
            avg_normal = np.mean(neighbor_normals, axis=0)
            curvature = 1 - np.abs(np.dot(normals[i], avg_normal))
            curvatures.append(curvature)
        
        curvatures = np.array(curvatures)
        
        # 基于曲率进行聚类分割
        features = np.column_stack([points, normals, curvatures.reshape(-1, 1)])
        clustering = DBSCAN(eps=0.15, min_samples=20).fit(features)
        labels = clustering.labels_
        
        # 为每个部件分配颜色
        max_label = labels.max()
        colors = plt.cm.tab10(np.linspace(0, 1, max_label + 1))[:, :3]
        
        # 噪声点设为灰色
        colored_labels = labels.copy()
        colored_labels[labels == -1] = max_label + 1
        colors = np.vstack([colors, [0.5, 0.5, 0.5]])
        
        pcd.colors = o3d.utility.Vector3dVector(colors[colored_labels])
        
        part_count = max_label + 1 if max_label >= 0 else 0
        noise_count = np.sum(labels == -1)
        print(f"曲率分割出 {part_count} 个部件，噪声点 {noise_count} 个")
        return pcd, labels